#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/8/13 19:07
# @Author  : CHEN Wang
# @Site    :
# @File    : timing_backtest.py
# @Software: PyCharm

"""
脚本说明: 单资产时间序列上的择时回测(输入指标，以及指标转换信号方法)
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
from quant_researcher.quant.datasource_fetch.index_api import index_price_related
from quant_researcher.quant.datasource_fetch.fund_api import fund_nav_related
from quant_researcher.quant.datasource_fetch.stock_api import stock_price_related
from quant_researcher.quant.datasource_fetch.factor_api import factor_exposure_related
from quant_researcher.quant.factors.factor_preprocess.preprocess import ts_percentile_rank_signal
from quant_researcher.quant.factors.indicator_tools.indicators_util import macd_signal, \
    kdj_signal, rsi_signal, wr_signal, bband_signal, ema_cross_signal, ma_cross_signal
from quant_researcher.quant.project_tool.db_operator import db_conn
from quant_researcher.quant.backtest.holding_backtest import simple_backtest
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.performance_attribution.core_functions.performance_analysis.performance import sharpe_ratio, annual_return, \
    max_drawdown, winning_rate, simple_perf
import warnings
from quant_researcher.quant.project_tool.exception import FOFServerError, ERR_S_2, FOFUserError, ERR_U_2

warnings.filterwarnings("ignore")
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


class TimingBacktest(object):
    def __init__(self, begin_date, end_date, asset_code, asset_type, indi_data=None,
                 indi_direction=1, commission=0, smoothing_window=0, **method_kwargs):
        """
        单一资产，通过指标转换成择时信号，进行择时策略回测

        :param str begin_date: 择时起始时间， 如'2020-01-01'
        :param str end_date: 择时终止时间， 如'2020-05-01'
        :param list asset_code: list of str, 进行择时的标的代码, 目前仅支持单一资产, eg. ['601900']
        :param str asset_type: 标的资产的类型, 支持股票(stock), 行业(industry), 基金(fund), 债券(bond), 指数(index)
        :param pd.DataFrame indi_data: 外部择时指标数据，需包含tradedate(格式需要为str)，和指标列; 目前仅支持单个指标
                                       默认为None, 即测试的指标为资产的收盘价相关指标
        一个包含ROE指标的数据形状如下：
        ———————————————————————
          |tradedate | roe
         |2010-01-03| 0.2
        |2010-01-03| 0.15
        ———————————————————————
        :param int indi_direction: 1或者-1, 表示因子的方向
        :param float commission: 手续费参数，默认为0
        :param int smoothing_window: 设置择时信号的平滑区间，确保新的择时信号必须持续window时长，避免择时信号切换太频繁，默认为0
        :param method_kwargs:
            - timing_method: func, 择时方法，支持timing_signal中所有方法，包括: macd,kdj,rsi,wr,bband,ema_cross,ma_cross等等, 默认macd
            - timing_params: dict, 对所选择的择时方法规定入参，每种择时方法的入参情况参见timing_signal, 默认为空, 即采用该种方法的默认参数
        """
        self.con = db_conn.get_derivative_data_conn(by_sqlalchemy=True)
        self.con_factor = db_conn.get_tk_factors_conn(by_sqlalchemy=True)
        self.con_astk = db_conn.get_stock_conn(by_sqlalchemy=True)
        self.begin_date = begin_date
        self.end_date = end_date
        self.indi_direction = indi_direction
        self.commission = commission
        self.asset_type = asset_type
        self.asset_code = asset_code
        self.smoothing_window = smoothing_window
        # timing_supported_all为所有支持的择时方法，timing_supported_shrink为部分支持的择时方法（除股票资产择时支持全部择时方法，其他均选择timing_supported_shrink）
        self.timing_supported_all = [macd_signal, kdj_signal, rsi_signal, wr_signal, bband_signal,
                                     ema_cross_signal, ma_cross_signal, ts_percentile_rank_signal]
        self.timing_supported_shrink = [macd_signal, rsi_signal, bband_signal, ema_cross_signal,
                                        ma_cross_signal, ts_percentile_rank_signal]
        self.timing_method = method_kwargs.pop('timing_method', macd_signal)
        self.timing_params = method_kwargs.pop('timing_params', {})
        supported = self.timing_supported_all if asset_type == 'stock' else self.timing_supported_shrink
        assert self.timing_method in supported, f"您所选的择时方法目前暂不支持，请在{[x.__name__ for x in supported]} 中选择"
        assert isinstance(smoothing_window, int) and smoothing_window >= 0, "请确保择时信号平滑窗口参数为非负整数"
        self.quote_data = self.get_quote_data()
        LOG.info(f"获取{asset_type}:{self.asset_code}的择时信号中")
        self.quote_data.rename(columns={'fap_close': 'close'}, inplace=True)
        self.quote_data['tradedate'] = pd.to_datetime(self.quote_data['tradedate'].astype(str))
        # 若择时信号中断，则延续之前的信号；若由于择时方法的计算方式，一开始有NA，则drop；
        if indi_data is None:
            self.indi_data = self.quote_data.set_index('tradedate').drop(
                columns=f"{self.asset_type}code")
        else:
            factor_name = indi_data.columns.difference(['tradedate']).tolist()[0]
            self.indi_data = indi_data.rename(columns={f'{factor_name}': 'close'})
            self.indi_data = self.indi_data[['close']]  # 取出dataframe形式
        print('择时指标和行情数据初始化完成')

    def update_indi(self, new_indi_data, new_direction, new_method):
        factor = new_indi_data.columns.difference(['tradedate']).tolist()[0]
        self.indi_data = new_indi_data.rename(columns={f'{factor}': 'close'})
        self.indi_data = self.indi_data.set_index('tradedate')[['close']]  # 取出dataframe形式
        self.timing_method = new_method
        self.indi_direction = new_direction
        LOG.info(f"获取{self.asset_type}:{self.asset_code}的择时信号中")
        self.signal = self.get_signal()

    def get_signal(self):
        """
        将self.indi_data转换为择时信号，并获取调仓日期

        :return:
        """

        signal = self.timing_method(self.indi_data, **self.timing_params).dropna()
        signal[signal != 1] = 0

        self.signal = signal
        if self.smoothing_window != 0:
            self.signal_smoothing()
        if self.indi_direction == -1:
            self.signal = - self.signal

    def get_trading_date(self):
        """
        根据self.signal择时信号获取调仓日期

        :return:
        """
        self.signal_trade = self.signal[
            self.signal != self.signal.shift(1).fillna(method='bfill')].dropna().index.tolist()

    def get_quote_data(self):
        if self.asset_type == 'stock':
            quote_data = stock_price_related.get_stock_quote(self.asset_code, self.begin_date, self.end_date,
                                                             select=['stock_code', 'end_date',
                                                                        'high', 'low', 'close as fap_close'])
            quote_data = quote_data.rename(columns={'stock_code': 'stockcode', 'end_date': 'tradedate'})
        elif self.asset_type == 'fund':
            # quote_data = fund_nav_related.get_fund_adj_nav(self.asset_code, [('nav_date', self.begin_date, self.end_date)],
            #                                                return_only_nav=False)
            quote_data = fund_nav_related.get_fund_nav(fund_code=self.asset_code, nav_type='adj',
                                                       start_date=self.begin_date, end_date=self.end_date)
            quote_data['end_date'] = pd.to_datetime(quote_data['end_date'])
            quote_data['end_date'] = quote_data['end_date'].dt.date
            quote_data = quote_data.rename(columns={'end_date': 'tradedate', 'fund_code': 'fundcode', 'adjusted_nav': 'fap_close'})
        elif self.asset_type == 'bond':
            quote_data = index_price_related.get_bond_index_quote(self.asset_code, self.begin_date, self.end_date,
                                                                  select=['trade_date as tradedate',
                                                                          'index_code as bondcode',
                                                                          'cfzs_val as fap_close'])
            quote_data['bondcode'] = [code.replace('.CS', '') for code in quote_data['bondcode']]
        elif self.asset_type == 'industry':
            quote_data = index_price_related.get_stock_index_quote(self.asset_code, self.begin_date, self.end_date,
                                                                   select=['trade_date as tradedate',
                                                                           'index_code as industrycode',
                                                                           'close as fap_close', 'high', 'low'])
        elif self.asset_type == 'index':
            quote_data = index_price_related.get_index_close(index_code=self.asset_code,
                                                             start_date=self.begin_date, end_date=self.end_date)
            if quote_data is None:
                LOG.error(f'无法获取该指数{self.asset_code}的行情数据，请检查并联系数据管理员')
                raise FOFServerError(ERR_S_2, f'无法获取该指数{self.asset_code}的行情数据，请检查并联系数据管理员')

            quote_data = quote_data.rename(columns={'end_date': 'trade_date'})
            quote_data = quote_data.rename(columns={'trade_date': 'tradedate', 'index_code': 'indexcode'})
        return quote_data

    def signal_smoothing(self):
        # 平滑择时信号
        self.get_trading_date()  # 获取调仓日期
        sig_t = [t for t in [self.signal.index[0]] + self.signal_trade]
        # 求每次信号的持续时间长度
        sig_len = [[sig_t[i], len(self.signal.loc[sig_t[i]:sig_t[i + 1], ]) - 1] for i in
                   range(len(sig_t) - 1)]
        for i in range(len(sig_len) - 1):
            if sig_len[i][1] <= self.smoothing_window:
                a_index = self.signal.index.tolist().index(sig_len[i + 1][0])
                self.signal.loc[sig_len[i][0]:self.signal.index[a_index - 1], ] = np.nan
        self.signal = self.signal.ffill()

    def get_plot(self, result):
        """
        画出择时策略净值曲线和超额收益曲线图

        :param result: DataFrame, 包含两个净值序列，一列为Portfolio，一列为Benchmark，索引为tradedate
        :return:
        """
        result = result.ffill()
        active = result.eval('Portfolio-Benchmark')
        fig, ax = plt.subplots(2, 1, figsize=(8, 9))
        title = '策略净值曲线'
        result.plot(ax=ax[0], grid=True, title=title, sharex=True, fontsize=14, figsize=(16, 9))
        active.plot(ax=ax[1], grid=True, title=f'超额收益', sharex=True, fontsize=14, figsize=(16, 9))
        ax[0].yaxis.set_major_formatter(mtick.PercentFormatter(1))
        ax[1].yaxis.set_major_formatter(mtick.PercentFormatter(1))
        # fig_path = os.path.join(BASE_DIR, f'{self.asset_code[0]}择时回测.png')
        # plt.savefig(fig_path)
        # plt.savefig(f'C:\\Users\\admin\\Desktop\\work\\files\\择时策略\\{self.asset_code[0]}择时回测.png')
        plt.show()

    def run(self, signal=None, plot=False):
        """
        主要的策略函数

        :param pd.Series signal: 1 和 -1 的择时信号；为外部直接指定的择时信号
        :param bool plot: 是否画出择时效果图
        :return:
        """
        print('开始择时策略的回测')
        if signal is None:  #未直接输入择时信号，则根据self.indi_data计算得到择时信号
            self.get_signal()  # 将self.indi_data转换为择时
        else:               # 直接输入了择时信号
            self.signal = signal
            self.signal_smoothing()

        asset_close_price = self.quote_data.set_index(['tradedate', f'{self.asset_type}code'])['close'].unstack().ffill()
        self.ret = asset_close_price.pct_change()
        self.asset_data = pd.concat([self.ret, self.signal], axis=1, keys=['ret', 'signal'], join='inner')
        self.asset_data.index = pd.to_datetime(self.asset_data.index)
        self.signal = self.signal.reindex(self.asset_data.index)

        # 构造simple_backtest所需的入参
        backtest_df = self.signal.to_frame()
        backtest_df.index = backtest_df.index.strftime('%Y-%m-%d').rename('tradedate')
        backtest_df.reset_index(inplace=True)
        backtest_df.rename(columns={0: 'weights'}, inplace=True)
        backtest_df['weights'] = backtest_df['weights'].replace(-1, 0)  # 择时信号1和-1转换为100% 和 0%的仓位
        backtest_df['code'] = self.asset_code

        portfolio_nav, benchmark_nav, portfolio_ret, perf = simple_backtest(
            asset_weights_df=backtest_df,
            asset_type=self.asset_type,
            benchmark=self.asset_code,
            benchmark_type=self.asset_type)

        # 加入手续费
        self.get_trading_date()  # 得到调仓日期
        commission_multiplier_dict = {date: (1 - self.commission) ** (i + 1) for i, date in enumerate(self.signal_trade)}
        commission_multiplier = pd.Series(commission_multiplier_dict)
        commission_multiplier = commission_multiplier.reindex(
            commission_multiplier.index.union(portfolio_nav.index)).ffill().reindex(portfolio_nav.index).fillna(1)
        self.commission_multiplier = commission_multiplier
        portfolio_nav = (portfolio_nav * commission_multiplier).rename('Portfolio')

        # 胜率
        self.asset_data = pd.concat([self.asset_data, portfolio_nav], axis=1, join='inner')
        self.asset_data.columns = ['ret', 'signal', 'cumret']
        win_rate_analysis = winning_rate(self.asset_data)

        nav_df = pd.DataFrame([portfolio_nav, benchmark_nav]).T
        nav_df.index = pd.to_datetime(nav_df.index)
        self.nav_df = nav_df

        if plot:
            self.get_plot(nav_df)

        self.perf = perf

        return self.signal, self.nav_df, asset_close_price


if __name__ == '__main__':
    begin_date = '2014-01-01'
    end_date = '2020-05-01'
    stock_code = ['601900']
    bond_code = ['CBA02501.CS']
    indu_code = ['801010']
    fund_code = ['007802']
    index_code = ['000300']

    # %% 测试收盘价相关的衍生指标择时效果
    # test1 = TimingBacktest(begin_date, end_date, stock_code, 'stock', indi_data=None, commission=0,
    #                        timing_method=ma_cross_signal, smoothing_window=10)
    # test2 = TimingBacktest(begin_date, end_date, bond_code, 'bond', indi_data=None, commission=0,
    #                        timing_method=macd_signal)
    # test3 = TimingBacktest(begin_date, end_date, indu_code, 'industry', indi_data=None,
    #                        commission=0, timing_method=ma_cross_signal)
    # test4 = TimingBacktest(begin_date, end_date, fund_code, 'fund', indi_data=None, commission=0,
    #                        timing_method=ma_cross_signal)
    # test5 = TimingBacktest(begin_date, end_date, index_code, 'index', indi_data=None, commission=0,
    #                        timing_method=ma_cross_signal)

    # %% 测试外部输入的衍生指标择时效果
    factor = 'epibs'
    con_factor = db_conn.get_tk_factors_conn(by_sqlalchemy=True)
    factor_table_name = pd.read_sql(f"select factor_table_name from stk_fac_info "
                                    f"where factor_name_en = '{factor}'", con_factor).iloc[0, 0]
    # factor_data = factor_exposure_related.get_stock_factor_exposure(factor, stock_code, begin_date, end_date)
    # factor_data.dropna(inplace=True)
    factor_data = factor_exposure_related.get_stock_factor_exposure(factor, stock_code, begin_date, end_date)
    factor_data = factor_data.rename(columns={'end_date': 'tradedate', 'stock_code': 'stockcode'})
    factor_data.dropna(inplace=True)
    factor_data['tradedate'] = pd.to_datetime(factor_data['tradedate'])
    indi_data = factor_data.where(factor_data.stockcode == stock_code[0]).dropna()
    # path = os.path.join(BASE_DIR, f'indi_data.h5')
    # indi_data.to_hdf(path, 'df')
    # indi_data = pd.read_hdf(path,'df')
    test6 = TimingBacktest(begin_date, end_date, stock_code, 'stock', indi_data=indi_data,
                           indi_direction=1, commission=0, timing_method=ma_cross_signal, plot=True)
    signal, _ = test6.run(test6.signal)
